import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import pickle
from training import misc
from training import dnnlib

# 设置随机种子以确保结果的可重复性
misc.set_seed(0)

# 加载预训练的StyleGAN模型
# StyleGAN2 FFHQ的预训练模型
url_ffhq = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-ffhq-config-f.pkl'
with dnnlib.util.open_url(url_ffhq, cache_dir='~/.cache/torch/hub') as f:
    G = pickle.load(f)

# 获取模型的w平均值（latent mean），用于生成图片
G = G.clone().to('cuda')  # 将模型移动到GPU
G.eval()  # 设置为评估模式
w_avg = G.mean_latent(4096)

# 生成随机噪声种子
seed = np.random.RandomState(0)
z = np.random.randn(1, G.z_dim)  # 随机采样一个潜在空间向量

# 将噪声种子转换为Tensor
z = tf.convert_to_tensor(z, dtype=tf.float32)

# 生成图片
images = G.components.synthesis.run(
    z, None, truncation_psi=0.7, randomize_noise=False)

# 将Tensor转换为NumPy数组，并进行后处理
images = tf.nn.tanh(images).numpy()
images = (images + 1) / 2.0  # 反标准化到[0, 1]

# 显示生成的图片

plt.imshow(images[0])
plt.axis('off')
plt.show()
